// Copyright (c) 2025 Shenzhen Kaihong Digital Industry Development Co., Ltd.
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use ohmw_core::bincode::{Decode, DefaultDecoder, DefaultEncoder, Encode};

pub struct MsgQueue {
    queue: Vec<Vec<u8>>,
}

impl MsgQueue {
    pub fn new() -> Self {
        Self { queue: Vec::new() }
    }

    pub fn push_bytes(&mut self, mut bytes: &[u8]) {
        while bytes.len() > 0 {
            if self.queue.last().is_none_or(|last| Self::bytes_is_complete(last.as_ref())) {
                if bytes.len() >= size_of::<u32>() {
                    let mut decoder = DefaultDecoder::new();
                    let data_size = unsafe { u32::decode_unchecked(&mut decoder, bytes) };
                    let all_size = size_of::<u32>() + data_size as usize;
                    bytes = &bytes[size_of::<u32>()..];

                    let mut encoder = DefaultEncoder::new();
                    let mut new_bytes = Vec::with_capacity(all_size);
                    data_size.encode(&mut encoder, &mut new_bytes).unwrap();
                    new_bytes.extend_from_slice(&bytes[..std::cmp::min(bytes.len(), data_size as usize)]);
                    self.queue.push(new_bytes);

                    if bytes.len() <= data_size as usize {
                        return;
                    }

                    bytes = &bytes[data_size as usize..];
                } else {
                    self.queue.push(Vec::from(bytes));
                    return;
                }
            } else {
                let last = self.queue.last_mut().unwrap();
                if last.len() < size_of::<u32>() {
                    if bytes.len() + last.len() < size_of::<u32>() {
                        last.extend_from_slice(bytes);
                        return;
                    }
                    
                    let remain_header_size = size_of::<u32>() - last.len();
                    last.extend_from_slice(&bytes[..remain_header_size]);
                    bytes = &bytes[remain_header_size..];
                    
                    let mut decoder = DefaultDecoder::new();
                    let data_size = unsafe { u32::decode_unchecked(&mut decoder, last.as_ref()) };
                    last.reserve(data_size as usize);
                }

                let mut decoder = DefaultDecoder::new();
                let remain_size = unsafe { u32::decode_unchecked(&mut decoder, last.as_ref()) } as usize - (last.len() - size_of::<u32>());
                last.extend_from_slice(&bytes[..std::cmp::min(bytes.len(), remain_size)]);
                if bytes.len() <= remain_size {
                    return;
                }

                bytes = &bytes[remain_size..];
            }
        }
    }

    pub fn drain_completed(&mut self) -> std::vec::Drain<'_, Vec<u8>> {
        if self.queue.last().is_none_or(|last| Self::bytes_is_complete(last)) {
            self.queue.drain(..)
        } else {
            self.queue.drain(..self.queue.len() - 1)
        }
    }

    fn bytes_is_complete(bytes: &[u8]) -> bool {
        if bytes.len() < size_of::<u32>() {
            return false;
        }

        let mut decoder = DefaultDecoder::new();
        let data_size = unsafe { u32::decode_unchecked(&mut decoder, bytes) };
        bytes.len() == data_size as usize + size_of::<u32>()
    }
}
